# Copyright (c) 2008, RPI-ACM, Jesse Kempf <jessekempf@gmail.com>
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# - Redistributions of source code must retain the above copyright notice,
#     this list of conditions and the following disclaimer.
# - Redistributions in binary form must reproduce the above copyright
#     notice, this list of conditions and the following disclaimer in the
#     documentation and/or other materials provided with the distribution.
# - Neither the names of RPI-ACM nor the names of its contributors
#     may be used to endorse or promote products derived from this software
#     without specific prior written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
#  "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
#  TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
#  PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER
#  OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
#  EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
#  PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
#  PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
#  LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
#  NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""PyNIX Reliable Message-Oriented Channel.

The MessageChannel provides a reliable, message-based communication facility
for use between processes that share an ancestor.
"""
import base64
import errno
import sys
import os
import platform
import socket

import unixutils.passfd

from distutils.version import LooseVersion as Version

MENC = base64.urlsafe_b64encode
MDEC = base64.urlsafe_b64decode

DEBUG = False

#
# Linux's implemenation of AF_UNIX SOCK_STREAMs is broken with
# respect to sendmsg(2). The use of SOCK_SEQPACKET seems to have
# the proper semantics. We keep a table to enumerate which OS presents
# which semantics.
#
# We throw the following exception if the semantics haven't been determined:
class AppropriateSocketTypeUnknown(EnvironmentError):
	"""msgsock has not been taught about the current OS."""
	pass
#
# The table is a list of (`osname', `vers', `socktype') triples.
# `osname' must be an exact match (case notwithstanding).
# `vers' must be a string indicating the minimum OS version which should
#   use the given socket type.
#
# NB: The table must be kept sorted so that each osname's versions are
#   listed from highest to lowest.
#

# pylint: disable-msg=C0103
_SOCKET_TYPE_TABLE = [
	('freebsd', '', socket.SOCK_STREAM),
	('openbsd', '', socket.SOCK_STREAM),
	('netbsd', '', socket.SOCK_STREAM),
	('darwin', '', socket.SOCK_STREAM),
	('linux', '2.6.4', socket.SOCK_SEQPACKET),
]

_SOCKET_TYPE = None

_my_os_name = platform.system().lower()
_my_os_vers = platform.release()

for (_os_name, _os_vers, _socktype) in _SOCKET_TYPE_TABLE:
	if _my_os_name == _os_name:
		if not _os_vers or Version(_my_os_vers) >= Version(_os_vers):
			_SOCKET_TYPE = _socktype
			break
else:
	raise AppropriateSocketTypeUnknown("msgsock must be taught about "
			"%s. Contact py-unixutil's maintainer for assistance." % 
			((_my_os_name, _my_os_vers),)
			)

# pylint: enable-msg=C0103


# pylint: disable-msg=R0901
class ChannelError(IOError):
	"""The base exception for errors across the MessageChannel."""
	pass

class ChannelFormatError(ChannelError):
	"""Indicates an incorrectly formatted message was
	sent across the MessageChannel."""
	pass

class ChannelIOError(ChannelError):
	"""Indicates an I/O error occurred on MessageChannel sockets."""
	pass

class ChannelClosedError(ChannelIOError):
	"""Indicates the socket underlying the MessageChannel was closed."""
	pass

# pylint: enable-msg=R0901

class MessageChannel(object):
	"""A reliable message-oriented channel.

	A MessageChannel provides a local, message-oriented IPC mechanism
	between two related processes."""
	TRAILER = '|&|'
	HDRLEN = 8
	CHUNKINGSIZ = 1024

	@staticmethod
	def msg_encode(opcode, args, kwargs):
		"""Encode the opcode, args, and kwargs for transmission
		over the channel."""

		arglist = ','.join([MENC(str(arg)) for arg in args])

		kwarglist = ','.join(["%s$%s" % (MENC(str(k)), MENC(str(v)))
			for k,v in kwargs.iteritems()])

		payload = "%s|%s|%s%s" % (opcode, arglist,
			kwarglist, MessageChannel.TRAILER)
		msg = "|%08x|%s" % (len(payload), payload)

		return msg

	@staticmethod
	def msg_decode(msg):
		"""Decode the given message into an opcode, args, and kwargs."""
		(_, mlen, payload) = msg.split('|', 2)

		mlen = int(mlen, 16)

		assert payload[mlen - 3:] == MessageChannel.TRAILER

		(opcode, args, kwargs, _, _) = payload.split('|')

		# Reconstitute the message from the components
		if args:
			args = [MDEC(x) for x in args.split(',')]
		else:
			args = list()

		if kwargs:
			kwargs = dict(
				[(MDEC(k), MDEC(v)) for k,v in
					[x.split('$') for x in kwargs.split(',')]
				]
			)
		else:
			kwargs = dict()


		return (opcode, args, kwargs)

	def __init__(self):
		"""Create a MessageChannel."""
		(self.sock0, self.sock1) = socket.socketpair(socket.AF_UNIX,
														_SOCKET_TYPE)
		self.socket = None
		self.sockno = -1
		self.msgq = []

		self.state = "CLOSED"

	def __tune_socket(self):
		"""Internal method."""
		self.setblocking()
		self.sockno = self.socket.fileno()
		self.state = "OPEN"

	def set_peer0(self):
		"""Configure the channel for running in one peer's process."""
		self.socket = self.sock0
		self.sock1.close()
		self.__tune_socket()

	set_parent = set_peer0

	def set_peer1(self):
		"""Configure the channel for running in the other peer's process."""
		self.socket = self.sock1
		self.sock0.close()
		self.__tune_socket()

	set_child = set_peer1

	def setblocking(self, blocking=True):
		"""Set blocking I/O mode of the channel."""
		self.socket.setblocking(blocking)

	def fileno(self):
		"""Return the descriptor backing the MessageChannel."""
		return self.sockno

	def sendva(self, opcode, *args, **kwargs):
		"""Send a message with the provided opcode, arguments,
		and keyword arguments."""
		self.send(opcode, args, kwargs)

	def send(self, opcode, args, kwargs):
		"""Send a message with the given opcode, argument list, and
		keyword argument dictionary across the channel."""

		enc_msg = MessageChannel.msg_encode(opcode, args, kwargs)
		try:
			written = 0
			while written < len(enc_msg):
				written += os.write(self.sockno, enc_msg[written:])
		except OSError, exc:
			if exc.errno == errno.EPIPE:
				raise ChannelClosedError("MessageChannel unexpectedly closed.")
			elif exc.errno == errno.EBADF and self.state != 'OPEN':
				raise ChannelClosedError("Attempt to write to closed socket.")
			else:
				raise

		if DEBUG:
			print "[%d]: MessageChannel.send(): (%d, %s, %s)\t%r" % \
				(os.getpid(), opcode, args, kwargs, enc_msg)
			sys.stdout.flush()


	def passfd(self, fdesc):
		"""Push a file descriptor across the channel."""

		try:
			unixutils.passfd.passfd(self.fileno(), fdesc)
		except OSError, exc:
			if exc.errno == errno.EBADF and self.state != 'OPEN':
				raise ChannelClosedError("Attempt to write to closed socket.")

	def recv(self):
		"""Read a message from the channel."""
		if not self.msgq:
			success = False
			while not success:
				try:
					recvbuf = os.read(self.sockno, MessageChannel.CHUNKINGSIZ)
					success = True
				except OSError, exc:
					if exc.errno == errno.EINTR:
						pass
					elif exc.errno == errno.EBADF and self.state != 'OPEN':
						raise ChannelClosedError("Attempt to read from closed "
													"socket.")
					else:
						raise

			if recvbuf == '':
				raise ChannelClosedError("MessageChannel unexpectedly closed.")
			if not recvbuf.startswith('|'):
				raise ChannelError("Message %r does not start with %r" %
					(recvbuf, '|'))

			size = int(recvbuf.split('|', 2)[1], 16)
			if size + MessageChannel.HDRLEN > MessageChannel.CHUNKINGSIZ:
				read = len(recvbuf)
				while read < size:
					try:
						recvbuf += os.read(self.sockno, size)
						read += len(recvbuf)
					except OSError, exc:
						if exc.errno == errno.EINTR:
							pass
						elif exc.errno == errno.EBADF and self.state != 'OPEN':
							raise ChannelClosedError("Attempt to read from "
													"closed socket.")
						else:
							raise

			msgbuf = recvbuf.split(MessageChannel.TRAILER)
			del msgbuf[-1]
			self.msgq = [msg + MessageChannel.TRAILER for msg in msgbuf]

		encmsg = self.msgq.pop(0)
		msg = MessageChannel.msg_decode(encmsg)

		if DEBUG:
			print "[%d]: MessageChannel.read(): %s\t%r" % \
				(os.getpid(), msg, encmsg)
			sys.stdout.flush()
		return msg

	def grabfd(self):
		"""Accept and return a file descriptor pushed across the channel."""
		try:
			return unixutils.passfd.grabfd(self.fileno())
		except OSError, exc:
			if exc.errno == errno.EBADF and self.state != 'OPEN':
				raise ChannelClosedError("Attempt to read from closed socket.")

	def close(self):
		"""Close this message channel and free associated resources."""
		self.state = "CLOSED"
		try:
			self.socket.close()
		except OSError, exc:
			if exc.errno == errno.EBADF:
				raise ChannelClosedError("Socket already closed.")
